#include <torch/torch.h>
#include <iostream>
#include <limits> // 用于设置初始的最大和最小值

// 模拟的训练数据集：key 和 position
torch::Tensor keys = torch::tensor({0.5, 1.0, 1.5, 2.0, 2.5}, torch::dtype(torch::kFloat32));
torch::Tensor positions = torch::tensor({0, 1, 2, 3, 4}, torch::dtype(torch::kFloat32));

// 定义线性模型
struct LinearModel : torch::nn::Module {
    torch::nn::Linear linear{nullptr};

    LinearModel() {
        // 构造一个线性模型 y = a * x + b，输入为1，输出为1
        linear = register_module("linear", torch::nn::Linear(1, 1));
    }

    // 前向传播
    torch::Tensor forward(torch::Tensor x) {
        return linear(x);
    }
};

int main() {
    // 实例化线性模型
    auto model = std::make_shared<LinearModel>();

    // 定义优化器 (使用 SGD)
    torch::optim::SGD optimizer(model->parameters(), torch::optim::SGDOptions(0.01));

    // 定义损失函数 (MSELoss)
    torch::nn::MSELoss loss_fn;

    // 训练数据集需要从一维向二维转换
    auto x_train = keys.view({-1, 1});
    auto y_train = positions.view({-1, 1});

    // 初始化最大和最小误差变量
    float max_loss = std::numeric_limits<float>::lowest();  // 初始最大值为负无穷
    float min_loss = std::numeric_limits<float>::max();      // 初始最小值为正无穷

    // 训练模型
    int num_epochs = 1000;
    for (int epoch = 0; epoch < num_epochs; ++epoch) {
        // 前向传播：计算模型的预测值
        torch::Tensor prediction = model->forward(x_train);

        // 计算损失
        torch::Tensor loss = loss_fn(prediction, y_train);

        // 反向传播并更新模型参数
        optimizer.zero_grad();
        loss.backward();
        optimizer.step();

        // 获取当前的损失值
        float current_loss = loss.item<float>();

        // 更新最大和最小损失
        if (current_loss > max_loss) {
            max_loss = current_loss;
        }

        if (current_loss < min_loss) {
            min_loss = current_loss;
        }

        // 每 100 个 epoch 输出一次损失值
        if (epoch % 100 == 0) {
            std::cout << "Epoch [" << epoch << "/" << num_epochs << "], Loss: " << current_loss << std::endl;
        }
    }

    // 训练结束后输出最大和最小误差值
    std::cout << "Maximum Loss during training: " << max_loss << std::endl;
    std::cout << "Minimum Loss during training: " << min_loss << std::endl;


    // 最终的误差值
    torch::Tensor final_prediction = model->forward(x_train);
    torch::Tensor final_loss = loss_fn(final_prediction, y_train);
    std::cout << "Final Loss: " << final_loss.item<float>() << std::endl;


    return 0;
}
